import numpy
from pylab import date2num, num2date
import copy

from doas import spec

def mean(spectra):
    """
    Given an iterable of doas.spec.Spectrum objects, returns a new Spectrum 
    object which is the mean of those provided. The properties of the returned 
    spectrum relate to the input spectra in the following ways:
    
    filename -> None (the returned spectrum does not relate to a spectrum file)
    wavelengths -> same as input spectra (which must all have the same wavelengths)
    counts -> mean of input counts weighted by integration time and n_avg of input spectra
    capture_time -> mean of input capture times weighted by n_avg
    n_avg -> sum of n_avg of input spectra
    int_time -> mean of input int_times weighted by n_avg of input spectra
    """
    
    #first, a few sanity checks
    if not hasattr(spectra, '__iter__'):
        raise TypeError("spectra argument must be iterable.")
    
    if not len(spectra) > 0:
        raise ValueError("No spectra to average")
    
    if len(spectra) == 1:
        new_spec = copy.deepcopy(spectra[0])
        new_spec.filename = None
        return new_spec
    
    #check that all spectra have the same wavelengths and spectrometers
    wavelengths = spectra[0].wavelengths
    spectrometer_id = spectra[0].spectrometer_id
    for s in spectra[1:]:
        if not numpy.all(wavelengths == s.wavelengths):
            raise ValueError("Cannot average spectra with different wavelengths")
        
        if spectrometer_id != s.spectrometer_id:
            raise ValueError("Cannot average spectra from different spectrometers")
        
    n_avgs = numpy.array([s.n_avg for s in spectra])
    int_times = numpy.array([s.int_time for s in spectra])
    capture_times_sec = numpy.array([date2num(s.capture_time) for s in spectra])
    
    #calculate the mean capture time weighted by n_avg
    new_capture_time = num2date(numpy.average(capture_times_sec, 0, n_avgs, False))
    
    #calculate the mean integration time weighted by the number of averages
    new_int_time, total_n_avg = numpy.average(int_times, 0, n_avgs, True)
    
    #stack the spectra counts together into an array
    all_counts = numpy.vstack([s.counts for s in spectra])
    
    #normalise them w.r.t. integration time
    norm_counts = all_counts.T * (new_int_time / int_times)
    
    #average them weighted by the number of averages
    new_counts = numpy.average(norm_counts, 1, n_avgs, False)
    
    #sanity check
    assert ((new_counts.shape == wavelengths.shape) or (len(wavelengths)== 0))
    
    new_spectrum = spec.Spectrum(None, wavelengths, new_counts, new_capture_time, 
                                 spectrometer_id, new_int_time, total_n_avg)

    return new_spectrum



def subtract(spec1, spec2):
    """
    Returns a new spectrum whose counts are equal to that of spec1 - spec2.
    If spec2 has a different integration time to spec1, then it is scaled to 
    that of spec1 prior to subtraction.
    """
    return spec._subtract(spec1, spec2)


class WavelengthRatioCalc:
    
    def __init__(self, wavelength1, wavelength2):
        
        self.__numerator = wavelength1
        self.__denominator = wavelength2
        self.__idx1 = None
        self.__idx2 = None
        self.__idx3 = None
        self.__idx4 = None
        self.__wavelengths = None
    
    def get_ratio(self, spec):
        if spec.wavelengths[0] > min(self.__numerator, self.__denominator):
            raise ValueError("Requested wavelength is outside of the range of the spectrum.")
        
        if spec.wavelengths[-1] < max(self.__numerator, self.__denominator):
            raise ValueError("Requested wavelength is outside of the range of the spectrum.")
    
        if self.__wavelengths is None or not numpy.all(spec.wavelengths == self.__wavelengths):
            #then we need to calculate all this stuff
            self.__wavelengths = spec.wavelengths.copy()
            
            idx = numpy.argmin(numpy.abs(self.__wavelengths - self.__numerator))

            if(self.__wavelengths[idx] > self.__numerator):
                self.__idx1 = idx - 1
                self.__idx2 = idx
            else:
                self.__idx1 = idx
                self.__idx2 = idx + 1
            
            idx = numpy.argmin(numpy.abs(self.__wavelengths - self.__denominator))
            
            if(self.__wavelengths[idx] > self.__denominator):
                self.__idx3 = idx - 1
                self.__idx4 = idx
            else:
                self.__idx3 = idx
                self.__idx4 = idx + 1    
                  
        cnts1 = spec.counts[self.__idx1] + ((self.__numerator - self.__wavelengths[self.__idx1])*((spec.counts[self.__idx2] - spec.counts[self.__idx1])/(self.__wavelengths[self.__idx2] - self.__wavelengths[self.__idx1])))
        cnts2 = spec.counts[self.__idx3] + ((self.__denominator - self.__wavelengths[self.__idx3])*((spec.counts[self.__idx4] - spec.counts[self.__idx3])/(self.__wavelengths[self.__idx4] - self.__wavelengths[self.__idx3])))
        
        return cnts1/cnts2
        
        
        
        